import torch 
from torch import nn
import torch.nn.functional as F

class BranchNet(nn.Module):
    def __init__(self,num_channel=512,num_classes = 10):
        super(BranchNet,self).__init__()
        
        self.pool = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Linear(num_channel,num_classes)
    def forward(self,x):
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

        
def BranchNet64(**kwargs):

    return BranchNet(num_channel=64,**kwargs)
def BranchNet128(**kwargs):

    return BranchNet(num_channel=128,**kwargs)
def BranchNet256(**kwargs):

    return BranchNet(num_channel=256,**kwargs)
def BranchNet512(**kwargs):

    return BranchNet(num_channel=512,**kwargs)


if __name__ == '__main__':

    pass
    